package main

import (
	"fmt"
	"sort"
)

func main() {
	nums := []int{-1, 0, 1, 2, -1, -4}
	res := threeSumIndex(nums)
	fmt.Println(res)
}

// 难点：右指针应该在右边界还是左边界；又或者是否不用确定边界而是在其他情况处理重复的数据

// 三数之和
// https://leetcode.cn/problems/3sum/
// 排序后，固定一个数，然后左指针指向第一个不同的，右指针指向最后一个不同的，寻找左右指针和为固定数的负数
func threeSum(nums []int) [][]int {
	// 排序
	sort.Ints(nums)

	n := len(nums)
	res := [][]int{}

	for i := range nums {

		target := -nums[i]

		// 找到第一个不同的数
		l, r := i+1, n-1
		for nums[r-1] == nums[r] {
			r--
		}
		for nums[l] == nums[i] {
			l++
		}

		for l < r {
			v := nums[l] + nums[r]
			if v == target {
				res = append(res, []int{nums[i], nums[l], nums[r]})
				break
			}
			// 右指针左移动
			if v > target {
				for l < r {
					r--
					if nums[r] != nums[r-1] {
						break
					}
				}
			} else { // 左指针往右
				for l < r {
					l++
					if nums[l-1] != nums[l] {
						break
					}
				}
			}
		}
	}

	return res
}

func threeSumIndex(nums []int) [][]int {
	sort.Ints(nums)
	n := len(nums)
	var res [][]int
	for i := 0; i < n; i++ {
		if i != 0 && nums[i-1] == nums[i] {
			continue //去掉重复元素
		}

		need := -nums[i] //需要的数
		tmpRes := twoSumTool(nums, i+1, need)
		if tmpRes == nil {
			continue
		}
		for _, item := range tmpRes {
			res = append(res, append(item, nums[i]))
		}
	}
	return res
}

func twoSumTool(nums []int, start, target int) [][]int {
	n := len(nums)
	l := start
	r := n - 1

	var res [][]int
	for l < r {
		left := nums[l]
		right := nums[r]
		v := left + right
		if v > target {
			r--
		} else if v < target {
			l++
		} else {
			res = append(res, []int{nums[l], nums[r]})
			for l < r && nums[l] == left {
				l++
			}
			for l < r && nums[r] == right {
				r--
			}
		}
	}
	return res
}

var _ = `
/* 注意：调用这个函数之前一定要先给 nums 排序 */
vector<vector<int>> nSumTarget(
    vector<int>& nums, int n, int start, int target) {

    int sz = nums.size();
    vector<vector<int>> res;
    // 至少是 2Sum，且数组大小不应该小于 n
    if (n < 2 || sz < n) return res;
    // 2Sum 是 base case
    if (n == 2) {
        // 双指针那一套操作
        int lo = start, hi = sz - 1;
        while (lo < hi) {
            int sum = nums[lo] + nums[hi];
            int left = nums[lo], right = nums[hi];
            if (sum < target) {
                while (lo < hi && nums[lo] == left) lo++;
            } else if (sum > target) {
                while (lo < hi && nums[hi] == right) hi--;
            } else {
                res.push_back({left, right});
                while (lo < hi && nums[lo] == left) lo++;
                while (lo < hi && nums[hi] == right) hi--;
            }
        }
    } else {
        // n > 2 时，递归计算 (n-1)Sum 的结果
        for (int i = start; i < sz; i++) {
            vector<vector<int>> 
                sub = nSumTarget(nums, n - 1, i + 1, target - nums[i]);
            for (vector<int>& arr : sub) {
                // (n-1)Sum 加上 nums[i] 就是 nSum
                arr.push_back(nums[i]);
                res.push_back(arr);
            }
            while (i < sz - 1 && nums[i] == nums[i + 1]) i++;
        }
    }
    return res;
}
`
